# 用于SSFM光纤每段仿真后的EDFA放大,并加入噪声
#   sig：输入需要放大的信号，复值形式,列向量或者数组
#   h:普朗克常数。
#   polar_n:偏振数,用于噪声功率的分配,双偏振情况下功率减半。
#   nf_db:噪声系数，dBm单位
#   B:信号带宽，Hz单位
#   f_carrier:被放大信号的载波频率,Hz
#   alpha_db:衰减系数，5dB/km单位
#   span_len:span长度km
import numpy as np


def amplifier(sig, EDFA_nf:int = 5, sampling_rate:float = 160e9, f_c:float = 193548387096774.16, gain:float = 10.0):
    """
    EDFA放大器
    :param sig: 信号 
    :param EDFA_nf: 信号 
    :param sampling_rate: 采样率 
    :param f_c: 信号 
    :param gain: 信号 
    :returns: 
    """
    prev_symbols = [sig[1],sig[2]]
    sig = sig[0]
    h = 6.626068e-34
    # polar_n = Signal['polarization']
    polar_n = 2
    # nf_db = Amplifier_para['EDFA_nf_db']
    nf_db = EDFA_nf

    # db to unit-less
    nf = 10 ** (nf_db / 10)

    # B = Signal['Sampling_rate']
    # f_carrier = Signal['f_c']
    B = sampling_rate
    f_carrier = f_c
    # alpha_db = Fiber_para['alpha_dB']
    # span_len = Fiber_para['L']
    # gain_db = alpha_db * span_len
    gain_db=gain

    # db to unit-less
    gain = 10 ** (gain_db / 10)

    data = sig

    # 信号放大
    # unit-less to Np
    g = np.log(gain)

    # 计算噪声功率
    n_pdf = h * f_carrier * (gain - 1) * nf
    n_power = n_pdf * B

    # 对单偏振或者双偏振添加噪声
    if polar_n == 2:
        #  双偏振前况
        data = np.array(data)
        data_x = data[0:, 0]
        data_y = data[0:, 1]
        data_x = data_x * np.exp(g / 2)
        data_y = data_y * np.exp(g / 2)

        # 噪声
        noise_ASE_real_x = np.random.randn(len(data_x))
        noise_ASE_imag_x = np.random.randn(len(data_x))
        n_x = np.sqrt(n_power / 2) * (noise_ASE_real_x + 1j * noise_ASE_imag_x)

        noise_ASE_real_y = np.random.randn(len(data_x))
        noise_ASE_imag_y = np.random.randn(len(data_x))
        n_y = np.sqrt(n_power / 2) * (noise_ASE_real_y + 1j * noise_ASE_imag_y)

        data_x = data_x + n_x
        data_y = data_y + n_y

        data_x = data_x.reshape(-1, 1)
        data_y = data_y.reshape(-1, 1)
        signal = np.concatenate((data_x, data_y), axis=1)
        print('补偿后功率')
        print(np.mean(np.abs(signal[0:,0]) ** 2))
    else:
        data = data * np.exp(g / 2)

        # Noise
        noise_ASE_real = np.random.randn(len(data))
        noise_ASE_imag = np.random.randn(len(data))
        n = np.sqrt(n_power) * (noise_ASE_real + 1j * noise_ASE_imag)

        signal = data + n

    return [signal, prev_symbols[0], prev_symbols[1]]

def Amplifier(sig, Amplifier_para, Signal, Fiber_para):
    h = 6.626068e-34
    polar_n = Signal['polarization']
    print(f'polar_n{polar_n}')
    nf_db = Amplifier_para['EDFA_nf_db']
    print(f'EDFA_nf_db{nf_db}')
    # db to unit-less
    nf = 10 ** (nf_db / 10)

    B = Signal['Sampling_rate']
    print(f'Sampling_rate{B}')
    f_carrier = Signal['f_c']
    print(f'f_c{f_carrier}')
    alpha_db = Fiber_para['alpha_dB']
    print(f'alpha_db{alpha_db}')
    span_len = Fiber_para['L']
    print(f'span_len{span_len}')
    gain_db = alpha_db * span_len
    print(f'gain_db{gain_db}')

    # db to unit-less
    gain = 10 ** (gain_db / 10)

    data = sig

    # 信号放大
    # unit-less to Np
    g = np.log(gain)

    # 计算噪声功率
    n_pdf = h * f_carrier * (gain - 1) * nf
    n_power = n_pdf * B

    # 对单偏振或者双偏振添加噪声
    if polar_n == 2:
        #  双偏振前况
        data = np.array(data)
        data_x = data[0:, 0]
        data_y = data[0:, 1]
        data_x = data_x * np.exp(g / 2)
        data_y = data_y * np.exp(g / 2)

        # 噪声
        noise_ASE_real_x = np.random.randn(len(data_x))
        noise_ASE_imag_x = np.random.randn(len(data_x))
        n_x = np.sqrt(n_power / 2) * (noise_ASE_real_x + 1j * noise_ASE_imag_x)

        noise_ASE_real_y = np.random.randn(len(data_x))
        noise_ASE_imag_y = np.random.randn(len(data_x))
        n_y = np.sqrt(n_power / 2) * (noise_ASE_real_y + 1j * noise_ASE_imag_y)

        data_x = data_x + n_x
        data_y = data_y + n_y

        data_x = data_x.reshape(-1, 1)
        data_y = data_y.reshape(-1, 1)
        signal = np.concatenate((data_x, data_y), axis=1)
        print('补偿后功率')
        print(np.mean(np.abs(signal[0:,0]) ** 2))
    else:
        data = data * np.exp(g / 2)

        # Noise
        noise_ASE_real = np.random.randn(len(data))
        noise_ASE_imag = np.random.randn(len(data))
        n = np.sqrt(n_power) * (noise_ASE_real + 1j * noise_ASE_imag)

        signal = data + n

    return signal
